#include "G4SystemOfUnits.hh"
#include "Randomize.hh"
#include "G4ParticleTable.hh"
#include "G4ParticleDefinition.hh"
#include <memory>

#include "MyGenCreator.hh"
#include "MyGenContainer.hh"

#include "CosmicMuon.hh"

CosmicMuon::CosmicMuon(): MyBaseGen()
{
    m_PGType = "Cosmic ray -- Muon";
    m_ClassName = "CosmicMuon";
    m_TypeFlag = 1005;
    m_EvtFlag = 0;

    m_Particle = "mu-";
    m_ParticleEnergy = 1. * MeV;
    m_ParticlePosition = G4ThreeVector(0., 0., 0.);
    m_ParticlePolarization = G4ThreeVector(0., 0., 0.);
    m_ParticleMomentumDirection = G4ThreeVector(0., 0., 0.);

    m_BoxLength = 30;
    m_BallRadius = 500;

    m_Gun = nullptr;
}

CosmicMuon::~CosmicMuon()
{
    if(!m_Gun)
        delete m_Gun;
}

void CosmicMuon::Initialize(YAML::Node yamlNode)
{
    m_Gun = new G4ParticleGun(1);

    m_Emin = yamlNode["Emin"].as<double>(1.) * MeV;
    m_Emax = yamlNode["Emax"].as<double>(1000.) * MeV;

    G4ParticleTable *particleTable = G4ParticleTable::GetParticleTable();
    m_mum = particleTable->FindParticle("mu-");
    m_mup = particleTable->FindParticle("mu+");

    m_Gun->SetParticleDefinition(m_mum);
    m_Gun->SetParticleEnergy(m_ParticleEnergy);
    m_Gun->SetParticlePolarization(m_ParticlePolarization);
}

void CosmicMuon::GeneratePrimaryVertex(G4Event* anEvent)
{
    // Generating
    G4double type = G4UniformRand();

    G4double energy = G4UniformRand() * (m_Emax - m_Emin) + m_Emin;
    G4double theta = G4UniformRand() * M_PI * 0.5;
    while(G4UniformRand() > pow(cos(theta), 2.5))
    {
       theta = G4UniformRand() * M_PI * 0.5;
    }
    G4double phi = 2 * G4UniformRand() * M_PI;
    G4double myX = m_BoxLength * (G4UniformRand() - 0.5);
    G4double myY = m_BoxLength * (G4UniformRand() - 0.5);

    std::pair<G4ThreeVector, G4ThreeVector> PosAndMomDir;

    PosAndMomDir.first = G4ThreeVector(-myX * sin(phi) + myY * cos(theta) * cos(phi) + m_BallRadius * sin(theta) * cos(phi),
                                      myX * cos(phi) + myY * cos(theta) * sin(phi) + m_BallRadius * sin(theta) * sin(phi),
                                      -myY * sin(theta) + m_BallRadius * cos(theta)) * cm;
    PosAndMomDir.second = G4ThreeVector(-sin(theta) * cos(phi), -sin(theta) * sin(phi), -cos(theta));

    m_ParticleEnergy = energy;
    m_ParticlePosition = PosAndMomDir.first;
    m_ParticleMomentumDirection = PosAndMomDir.second;

    // Setting
    if (type < 0.5) m_Gun->SetParticleDefinition(m_mum);
    else m_Gun->SetParticleDefinition(m_mup);
    m_Gun->SetParticleEnergy(energy);
    m_Gun->SetParticlePosition(m_ParticlePosition);
    m_Gun->SetParticleMomentumDirection(m_ParticleMomentumDirection);

    m_Gun->GeneratePrimaryVertex(anEvent);
    ++m_EvtFlag;
}

REGISTGPSGEN(CosmicMuon)